/**
 * @file encrypt.c
 * 
 * @date: 2021-08-05
 * 
 * @author: lyfan
 * 
 * library encrypt code
**/
#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <stdbool.h>
#include <stdint.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#include <errno.h>
#include <signal.h>
#include <string.h>
#include <getopt.h>
#include <openssl/des.h>

#define		DES_NONE_STR				"none"
#define		DES_ECB_STR					"ecb"
#define		DES_ECB3_STR				"ecb3"
#define		DES_CBC_STR					"cbc"
#define		CRYPTMAGIC					"WEIZ"

typedef enum _DES_CRYPT_TYPE{
	DES_NONE = 0,
	DES_ECB,
	DES_ECB3,
	DES_CBC,

	DES_MAX = 255,
} DES_CRYPT_TYPE;

char *cryptstring[] = {
	DES_NONE_STR,
	DES_ECB_STR,
	DES_ECB3_STR,
	DES_CBC_STR,
};

typedef struct __WX_CRYPT_HEAD
{
	char 			szMagic[4];
	uint32_t		filelen;
	uint8_t			cryptype;
	uint8_t			reserved[3];
} WX_CRYPT_HEAD;

static struct option long_options[] = {
	{"crypt", required_argument, NULL, 'c'},
	{"file", required_argument, NULL, 'f'},
	{"list", no_argument, NULL, 'l'},
	{"decrypt", no_argument, NULL, 'd'},
	{NULL, 0, NULL, 0}
};

static unsigned char key_data[6][8] = {
	{0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01},
	{0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01},
	{0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01},
	{0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01},
	{0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01},
	{0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01},
};

static unsigned char cbc_key[8] = {0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01};
static unsigned char cbc_iv[8] = {0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01};

void usage(void)
{
	fprintf(stderr,
		"Usage: encrypt [OPTIONS]... [ARGS]...\n"
		"en/decrypt library file, support ecb, ecb3, cbc, ...\n"
		"\n"
		" -c, --crypt [ecb, ecb3, cbc, ...]\n"
		" -f, --file path2file\n"
		" -d, --de decrypt, default encrypt\n"
		"-l, --list crypt method support"
		"\n");

	return ;
}

void support(void)
{
	fprintf(stderr,
		"encrypt support ecb, ecb3, cbc, ...\n\n");

	return ;
}

void SF_LogHex(uint8_t *buf, int len)
{
	int i = 0;

	while(i < len){
		fprintf(stdout, "%02x ", buf[i]);
		if ((i+1) % 8 == 0){
			fprintf(stdout, "\n");
		}
		i++;
	}
	fprintf(stdout, "\n");

	return ;
}

static int ssl_des_ecb(const char *datain, char *dataout, int crypt)
{
	DES_key_schedule ks;
	DES_cblock in, out;

	DES_set_key_unchecked(&key_data[0], &ks);
	memcpy(in, datain, 8);
	memset(out, 0, 8);
	memset(dataout, 0, 8);
	if (crypt == DES_ENCRYPT)
		DES_ecb_encrypt(&in, &out, &ks, DES_ENCRYPT);
	else
		DES_ecb_encrypt(&in, &out, &ks, DES_DECRYPT);
	memcpy(dataout, out, 8);

	return 0;
}

static int ssl_des_ede_ecb(const char *datain, char *dataout, int crypt)
{
	DES_cblock in, out;
	DES_key_schedule ks, ks2, ks3;

	DES_set_key_unchecked(&key_data[0], &ks);
	DES_set_key_unchecked(&key_data[1], &ks2);
	DES_set_key_unchecked(&key_data[2], &ks3);
	memcpy(in, datain, 8);
	memset(out, 0, 8);
	memset(dataout, 0, 8);
	if (crypt == DES_ENCRYPT)
		DES_ecb3_encrypt(&in, &out, &ks, &ks2, &ks, DES_ENCRYPT);
	else
		DES_ecb3_encrypt(&in, &out, &ks, &ks2, &ks, DES_DECRYPT);
	memcpy(dataout, out, 8);

	return 0;
}

static int ssl_des_cbc(const char *datain, char *dataout, int crypt)
{
	unsigned char cbc_in[8];
	unsigned char cbc_out[8];
	DES_cblock iv3;
	DES_key_schedule ks;

	DES_set_key_checked(&cbc_key, &ks);
	memset(cbc_out, 0, sizeof(cbc_out));
	memset(cbc_in, 0, sizeof(cbc_in));
	memset(dataout, 0, 8);
	memcpy(cbc_in, datain, 8);
	memcpy(iv3, cbc_iv, sizeof(cbc_iv));
	if (crypt == DES_ENCRYPT)
		DES_ncbc_encrypt(cbc_in, cbc_out, 8, &ks, &iv3, DES_ENCRYPT);
	else
		DES_ncbc_encrypt(cbc_in, cbc_out, 8, &ks, &iv3, DES_DECRYPT);
	memcpy(dataout, cbc_out, 8);

	return 0;
}

static int sslcrypt_lib_entry(const char *filein, char *fileout, int roundlen, DES_CRYPT_TYPE ctype, int crypt)
{
	int i = 0;

	for (; i < roundlen; i += 8){
		switch (ctype){
			case DES_ECB:
				ssl_des_ecb(filein, fileout, crypt);
				break;
			case DES_ECB3:
				ssl_des_ede_ecb(filein, fileout, crypt);
				break;
			case DES_CBC:
				ssl_des_cbc(filein, fileout, crypt);
				break;
			default:
				break;
		}
		filein += 8;
		fileout += 8;
	}

	return 0;
}

int main(int argc, char **argv)
{
	int c, fild = 0;
	int sofd, genfd, rlen, datalen;
	uint32_t filelen, roundlen, writelen = 0;
	struct stat statbuf;
	char soname[256] = {0}, gensoname[256] = {0}, cryptname[64] = {0}/*, *strpos = NULL */;
	char *filein = NULL, *fileout = NULL;
	DES_CRYPT_TYPE ctype = DES_NONE;
	bool encrypted = true;

	while ((c = getopt_long(argc, argv, "c:f:l:d", long_options, NULL)) != EOF){
		switch (c){
		case 'c':
			strncpy(cryptname, argv[optind-1], sizeof(cryptname));
			if (memcmp(DES_ECB_STR, cryptname, strlen(DES_ECB_STR)) == 0 && (strlen(cryptname) == strlen(DES_ECB_STR))){
				ctype = DES_ECB;
			}
			if (memcmp(DES_ECB3_STR, cryptname, strlen(DES_ECB3_STR)) == 0 && (strlen(cryptname) == strlen(DES_ECB3_STR))){
				ctype = DES_ECB3;
			}
			if (memcmp(DES_CBC_STR, cryptname, strlen(DES_CBC_STR)) == 0 && (strlen(cryptname) == strlen(DES_CBC_STR))){
				ctype = DES_CBC;
			}
			break;
		case 'f':
			fild = 1;
			strncpy(soname, argv[optind-1], sizeof(soname));
			break;
		case 'l':
			support();
			exit(0);
			break;
		case 'd':
			encrypted = false;
			break;
		default:
			fprintf(stderr, "Unknown args!\n");
			usage();
			exit(-1);
			break;
		}
	}

	if (fild == 0 || (encrypted && ctype == DES_NONE)){
		usage();
		if (encrypted && ctype == DES_NONE) support();
		exit(-1);
	}

	sofd = open(soname, O_RDONLY);
	if (sofd < 0){
		printf("open %s failed\n", soname);
		return -1;
	}

	fstat(sofd, &statbuf);
	filelen = statbuf.st_size;
	roundlen = (filelen/8) * 8 + (filelen % 8 ? 8 : 0);
	filein = malloc(roundlen);
	if (filein == NULL){
		printf("malloc filein failed, size: %d\n", roundlen);
		close(sofd);
		return -1;
	}

	memset(filein, 0, roundlen);
	datalen = 0;
	while (datalen < filelen){
		rlen = read(sofd, filein+datalen, 4096 <= (filelen - datalen) ? 4096 : (filelen - datalen));
		if (rlen < 0 && errno != EAGAIN && errno != EINTR){
			printf("read lib error: %s\n", strerror(errno));
			free(filein), close(sofd);
			return -1;
		}
		datalen += rlen;
	}

	if (encrypted){
		if (memcmp(filein, "\177ELF", 4) != 0){
			printf("lib broken; not ELF\n");
			free(filein), close(sofd);
			return -1;
		}
		fileout = malloc(roundlen + sizeof(WX_CRYPT_HEAD));
		if (fileout == NULL){
			printf("malloc fileout failed, size: %d\n", roundlen);
			free(filein); close(sofd);
			return -1;
		}

		memset(fileout, 0, roundlen + sizeof(WX_CRYPT_HEAD));
		switch (ctype){
			case DES_ECB:
				sslcrypt_lib_entry(filein, fileout + sizeof(WX_CRYPT_HEAD), roundlen, DES_ECB, DES_ENCRYPT);
				memcpy(fileout, CRYPTMAGIC, 4);
				((WX_CRYPT_HEAD *)fileout)->filelen = filelen;
				((WX_CRYPT_HEAD *)fileout)->cryptype = DES_ECB;
				break;
			case DES_ECB3:
				sslcrypt_lib_entry(filein, fileout + sizeof(WX_CRYPT_HEAD), roundlen, DES_ECB3, DES_ENCRYPT);
				memcpy(fileout, CRYPTMAGIC, 4);
				((WX_CRYPT_HEAD *)fileout)->filelen = filelen;
				((WX_CRYPT_HEAD *)fileout)->cryptype = DES_ECB3;
				break;
			case DES_CBC:
				sslcrypt_lib_entry(filein, fileout + sizeof(WX_CRYPT_HEAD), roundlen, DES_CBC, DES_ENCRYPT);
				memcpy(fileout, CRYPTMAGIC, 4);
				((WX_CRYPT_HEAD *)fileout)->filelen = filelen;
				((WX_CRYPT_HEAD *)fileout)->cryptype = DES_CBC;
				break;
			default:
				printf("unknown crypt type-%d\n", ctype);
				free(filein), free(fileout), close(sofd);
				exit(-1);
				break;
		}
		strcpy(gensoname, "./encrypt.so");
		writelen = filelen + sizeof(WX_CRYPT_HEAD);
	} else{		// decrypt
		if (memcmp(filein, CRYPTMAGIC, 4) != 0){
			printf("no crypt, exit!\n");
			free(filein), close(sofd);
			return -1;
		}

		fileout = malloc(roundlen - sizeof(WX_CRYPT_HEAD));
		if (fileout == NULL){
			printf("malloc fileout failed, size: %lu\n", roundlen - sizeof(WX_CRYPT_HEAD));
			free(filein); close(sofd);
			return -1;
		}

		memset(fileout, 0, roundlen - sizeof(WX_CRYPT_HEAD));
		switch (((WX_CRYPT_HEAD *)filein)->cryptype){
			case DES_ECB:
				sslcrypt_lib_entry(filein + sizeof(WX_CRYPT_HEAD), fileout, roundlen - sizeof(WX_CRYPT_HEAD), DES_ECB, DES_DECRYPT);
				break;
			case DES_ECB3:
				sslcrypt_lib_entry(filein + sizeof(WX_CRYPT_HEAD), fileout, roundlen - sizeof(WX_CRYPT_HEAD), DES_ECB3, DES_DECRYPT);
				break;
			case DES_CBC:
				sslcrypt_lib_entry(filein + sizeof(WX_CRYPT_HEAD), fileout, roundlen - sizeof(WX_CRYPT_HEAD), DES_CBC, DES_DECRYPT);
				break;
			default:
				printf("unknown crypt type-%d\n", ctype);
				free(filein), free(fileout), close(sofd);
				exit(-1);
				break;
		}
		strcpy(gensoname, "./decrypt.so");
		printf("recognize %s, decrypt...\n", cryptstring[((WX_CRYPT_HEAD *)filein)->cryptype]);
		writelen = ((WX_CRYPT_HEAD *)filein)->filelen;
	}

	genfd = open(gensoname, O_RDWR | O_CREAT | O_TRUNC);
	if (genfd < 0){
		printf("open %s failed\n", gensoname);
		free(filein), free(fileout), close(sofd);
		return -1;
	}

	datalen = 0;
	while(datalen < writelen){
		rlen = write(genfd, fileout+datalen, 4096 <= (writelen - datalen) ? 4096 : (writelen - datalen));
		if (rlen < 0 && errno != EAGAIN && errno != EINTR){
			printf("write error: %s\n", strerror(errno));
			free(filein), free(fileout), close(sofd), close(genfd);
			return -1;
		}
		datalen += rlen;
	}

	free(filein);
	free(fileout);
	close(sofd);
	close(genfd);

	return 0;
}



